#include "MVDataAgentXML.h"
#include "tinyxml.h"
#include <string>
#include "MVDataFrame.h"
#include "MVDataSubMB.h"
#include "MVDataSubSubMB.h"

const char * FrameTypeStringMap[] = { "P_SLICE", "I_SLICE", "B_SLICE", "UNKNOWN" };
const char * MBTypeStringMap[] = { "SKIP", "P16x16", "P16x8", "P8x16", "P8x8", "I4MB", "I16MB", "UNKNOWN" };
const char * SubMBTypeStringMap[] = { "SMB8x8", "SMB8x4", "SMB4x8", "SMB4x4", "UNKNOWN" };
const char * LocStringMap[] = { "up", "upright", "right", "downright", "down", "downleft", "left", "upleft", "full", "UNKNOWN" };

CMVDataAgentXML::CMVDataAgentXML( CMVData * mvdata, const std::string dataname )
{
	m_pMVData = mvdata;
	//m_strDataName = dataname;
}

CMVDataAgentXML::~CMVDataAgentXML()
{
	// nothing to do now
}

int CMVDataAgentXML::ReadMVData()
{
	TiXmlDocument * pMVFile = new TiXmlDocument( m_pMVData->GetDataName().c_str() );

	if( ! pMVFile->LoadFile() )
	{
		return -1;
	}

	// read the MV info from the file
	TiXmlNode* fileNode = pMVFile->FirstChild( "file" );
	TiXmlNode* node = NULL;

	int width;
	int height;
	TiXmlElement* fileelement = fileNode->ToElement();
	fileelement->QueryIntAttribute( "width", &width );
	fileelement->QueryIntAttribute( "height", &height );

	m_pMVData->m_iFrameWidth = width;
	m_pMVData->m_iFrameHeight = height;

	while( node = fileNode->IterateChildren( "frame", node ) ) 
	{
		TiXmlElement* frameelement = node->ToElement();
		CMVDataFrame * frame = _ReadFrame( frameelement );
		// TODO: add frame to list

		m_pMVData->AddFrame( frame );
	}

	delete pMVFile;

	return 0;
}


CMVDataFrame * CMVDataAgentXML::_ReadFrame( TiXmlElement * frameelement )
{

	int frame_no;
	TiXmlNode * node = NULL;
	frameelement->QueryIntAttribute( "no", &frame_no );
	CMVDataFrame * frame = new CMVDataFrame( frame_no, _ConvertStrToFrameType( frameelement->Attribute( "type" ) ) );

	while( node = frameelement->IterateChildren( "macroblock", node ) ) 
	{
		TiXmlElement* mbelement = node->ToElement();

		CMVDataMB * mb = _ReadMB( mbelement );

		frame->AddMB( mb );

	}

	return frame;
}

CMVDataMB * CMVDataAgentXML::_ReadMB( TiXmlElement * mb_ele )
{
	int i = 0;
	TiXmlNode* sub_n = NULL;

	int bx, by;
	mb_ele->QueryIntAttribute( "xx", &bx );
	mb_ele->QueryIntAttribute( "yy", &by );

	CMVDataMB * mb = new CMVDataMB( _ConvertStrToMBType( mb_ele->Attribute( "type" ) ), bx, by );

	int vx, vy;
	
	switch ( mb->GetMBType() ) 
	{
	case TYPE_MB_P16x16:
		_ReadPredMV( mb_ele->FirstChild( "mv" )->ToElement(), vx, vy );
		mb->SetMVCoordinate( vx, vy );
		break;

	case TYPE_MB_P16x8:
	case TYPE_MB_P8x16:
	case TYPE_MB_P8x8:
		sub_n = NULL;
		while ( sub_n = mb_ele->IterateChildren( "sub_mb", sub_n ) ) 
		{
			TiXmlElement* sub_mb_element = sub_n->ToElement();
			CMVDataSubMB * sub_mb = _ReadSubMB( sub_mb_element );
			mb->AddSubMB( sub_mb );
		}
		break;

	}

	return mb;
}

CMVDataSubMB * CMVDataAgentXML::_ReadSubMB( TiXmlElement * sub_mb_ele )
{
	CMVDataSubMB * sub_mb = new CMVDataSubMB();

	LOC_BLOCK loc = _ConvertStrToLocation( sub_mb_ele->Attribute( "where" ) );
	sub_mb->SetLoc( loc );
	
	const char * sub_type = sub_mb_ele->Attribute( "subtype" );
	if( NULL != sub_type )
	{
		sub_mb->SetSubMBType( _ConvertStrToSubMBType( sub_type ) );
	}

	TiXmlNode * pMVNode = sub_mb_ele->FirstChild( "mv" );
	if( NULL != pMVNode )
	{
		int vx, vy;
		_ReadPredMV( pMVNode->ToElement(), vx, vy );
		sub_mb->SetMVCoordinate( vx, vy );
	}
	else
	{
		TiXmlNode * sub_n = NULL;
		while ( sub_n = sub_mb_ele->IterateChildren( "sub_sub_mb", sub_n ) )
		{
			CMVDataSubSubMB * sub_sub_mb = _ReadSubSubMB( sub_n->ToElement() );
			sub_mb->AddSubSubMB( sub_sub_mb );
		}

	}

	return sub_mb;

}

CMVDataSubSubMB * CMVDataAgentXML::_ReadSubSubMB( TiXmlElement * sub_sub_mb_ele )
{

	CMVDataSubSubMB * sub_sub_mb = new CMVDataSubSubMB();
	int vx, vy;

	LOC_BLOCK loc = _ConvertStrToLocation( sub_sub_mb_ele->Attribute( "where" ) );

	sub_sub_mb->SetLoc( loc );

	_ReadPredMV( sub_sub_mb_ele->FirstChild( "mv" )->ToElement(), vx, vy );
	sub_sub_mb->SetMVCoordinate( vx, vy );


	return sub_sub_mb;
}

void CMVDataAgentXML::_ReadPredMV( TiXmlElement * pre_mv_ele, int& vx, int& vy )
{
	pre_mv_ele->QueryIntAttribute( "vx", &vx );
	pre_mv_ele->QueryIntAttribute( "vy", &vy );
	return;
}



int CMVDataAgentXML::WriteMVData()
{

	TiXmlDocument * pMVFile = new TiXmlDocument( m_pMVData->GetDataName().c_str() );

	TiXmlDeclaration * pXmlDeclaration = new TiXmlDeclaration( "1.0", "UTF-8", "" );
	pMVFile->InsertEndChild( *pXmlDeclaration );

	TiXmlElement * pFileEle = new TiXmlElement( "file" );
	pFileEle->SetAttribute( "name", m_pMVData->GetDataName().c_str() );
	pFileEle->SetAttribute( "width", m_pMVData->GetFrameWidth() );
	pFileEle->SetAttribute( "height", m_pMVData->GetFrameHeight() );

	const CMVDataFrame * pFrame = NULL;
	const CMVDataMB * pMB = NULL;
	const CMVDataSubMB * pSubMB = NULL;
	const CMVDataSubSubMB * pSubSubMB = NULL;

	while( pFrame = m_pMVData->GetNextFrame( pFrame ) )
	{

		TiXmlElement * pFrameEle = new TiXmlElement( "frame" );
		pFrameEle->SetAttribute( "no", pFrame->GetFrameNumber() );
		pFrameEle->SetAttribute( "type", FrameTypeStringMap[pFrame->GetFrameType()] );

		while( pMB = pFrame->GetNextMB( pMB ) )
		{

			TiXmlElement * pMBEle = new TiXmlElement( "macroblock" );
			int x, y;
			pMB->GetCoordinate( x, y );

			pMBEle->SetAttribute( "xx", x );
			pMBEle->SetAttribute( "yy", y );
			pMBEle->SetAttribute( "type", MBTypeStringMap[pMB->GetMBType()] );
			// pMBEle->SetAttribute( "ref", this_frame-1 ); //TODO add ref

			if( !pMB->HaveSubMB() )
			{
				if ( pMB->HaveMVCoordinate() )
				{
					TiXmlElement * pMVEle = new TiXmlElement( "mv" );
					int vx, vy;
					pMB->GetMVCoordinate( vx, vy );
					pMVEle->SetAttribute( "vx", vx );
					pMVEle->SetAttribute( "vy", vy );
					pMBEle->InsertEndChild( *pMVEle );
				}
			}
			else
			{

				while( pSubMB = pMB->GetNextSubMB( pSubMB ) )
				{
					TiXmlElement * pSubMBEle = new TiXmlElement( "sub_mb" );

					pSubMBEle->SetAttribute( "where", LocStringMap[pSubMB->GetLoc()] );
					TYPE_SUB_MB submb_type = pSubMB->GetSubMBType();
					if( submb_type != TYPE_SUB_MB_NO_USE)
					{
						pSubMBEle->SetAttribute( "subtype", SubMBTypeStringMap[submb_type] );
					}

					if( !pSubMB->HaveSubSubMB() )
					{
						if ( pSubMB->HaveMVCoordinate() )
						{
							TiXmlElement * pMVEle = new TiXmlElement( "mv" );
							int vx, vy;
							pSubMB->GetMVCoordinate( vx, vy );
							pMVEle->SetAttribute( "vx", vx );
							pMVEle->SetAttribute( "vy", vy );
							pSubMBEle->InsertEndChild( *pMVEle );
						}
					}
					else
					{
						while( pSubSubMB = pSubMB->GetNextSubSubMB( pSubSubMB ) )
						{
							
							TiXmlElement * pSubSubMBEle = new TiXmlElement( "sub_sub_mb" );

							pSubSubMBEle->SetAttribute( "where", LocStringMap[pSubSubMB->GetLoc()] );

							TiXmlElement * pMVEle = new TiXmlElement( "mv" );
							int vx, vy;
							pSubSubMB->GetMVCoordinate( vx, vy );
							pMVEle->SetAttribute( "vx", vx );
							pMVEle->SetAttribute( "vy", vy );
							pSubSubMBEle->InsertEndChild( *pMVEle );

							pSubMBEle->InsertEndChild( *pSubSubMBEle );


						} // while( pSubSubMB )

					} // if( pSubMB->m_bSetMV )

					pMBEle->InsertEndChild( *pSubMBEle );

				} // while( pSubMB )

			} // if( pMB->m_bSetMV )

			pFrameEle->InsertEndChild( *pMBEle );

		} // while( pMB )

		pFileEle->InsertEndChild( *pFrameEle );

	} // while( pFrame );

	pMVFile->InsertEndChild( *pFileEle );
	pMVFile->SaveFile();

	return 0;

}


TYPE_FRAME CMVDataAgentXML::_ConvertStrToFrameType( const char * frametype )
{
	for( unsigned int i=0; i<TYPE_FRAME_MAX; i++ )
	{
		if( strcmp( frametype, FrameTypeStringMap[i] ) == 0 )
		{
			return (TYPE_FRAME) i;
		}
	}

	return TYPE_FRAME_MAX;

}

LOC_BLOCK CMVDataAgentXML::_ConvertStrToLocation( const char * loc )
{
	for( unsigned int i=0; i<LOC_BLOCK_MAX; i++ )
	{
		if( strcmp( loc, LocStringMap[i] ) == 0 )
		{
			return (LOC_BLOCK) i;
		}
	}

	return LOC_BLOCK_MAX;

}

TYPE_MB CMVDataAgentXML::_ConvertStrToMBType( const char * mbtype )
{
	for( unsigned int i=0; i<TYPE_MB_MAX; i++ )
	{
		if( strcmp( mbtype, MBTypeStringMap[i] ) == 0 )
		{
			return (TYPE_MB) i;
		}
	}

	return TYPE_MB_NO_USE;

}

TYPE_SUB_MB CMVDataAgentXML::_ConvertStrToSubMBType( const char * submbtype )
{
	for( unsigned int i=0; i<TYPE_SUB_MB_MAX; i++ )
	{
		if( strcmp( submbtype, SubMBTypeStringMap[i] ) == 0 )
		{
			return (TYPE_SUB_MB) i;
		}
	}

	return TYPE_SUB_MB_NO_USE;

}

int CMVDataAgentXML::Close()
{
	return 0;
}